/*
* Copyright (C) 2012-2014 qar
* License: http://www.gnu.org/licenses/gpl.html GPL version 2
*/

#include "cipher.h"

#include "util.h"
#include "crypto/ciphers/common.h"
#include "crypto/ciphers/encryptor.h"
#include "crypto/ciphers/encryptornull.h"
#include "crypto/ciphers/encryptoraes.h"
#include "crypto/ciphers/encryptorsalsa20.h"

Cipher::Cipher(Cipher::EncryptionAlgorithm algorithm, QString keyFile)
    : algorithm(algorithm),
      keyFile(keyFile),
      encryptor(NULL)
{
	setAlgorithm(algorithm, keyFile, false);
}

Cipher::~Cipher()
{
    delete encryptor;
    encryptor = NULL;
}

void Cipher::setAlgorithm(Cipher::EncryptionAlgorithm algorithm, QString keyFile, bool canGenerate)
{
    delete encryptor;
    encryptor = NULL;

    this->algorithm = algorithm;
    this->keyFile = keyFile;

    if (algorithm == Cipher::NoEncryption) {
        encryptor = new EncryptorNull();
	} else if (algorithm == Cipher::Aes256) {
		QByteArray key = getKey();
		if (key.isEmpty()) {
			if (canGenerate) {
				fprintf(stderr, "Generating key in %s...\n", keyFile.toLatin1().constData());
				generateKey(4096);
			} else {
				fprintf(stderr, "Key not found in %s!\n", keyFile.toLatin1().constData());
			}
		} else {
			fprintf(stderr, "Loading key from %s...\n", keyFile.toLatin1().constData());
		}
		encryptor = new EncryptorAes(key);
	} else if (algorithm == Cipher::Salsa20) {
		QByteArray key = getKey();
		if (key.isEmpty()) {
			if (canGenerate) {
				fprintf(stderr, "Generating key in %s...\n", keyFile.toLatin1().constData());
				generateKey(4096);
			} else {
				fprintf(stderr, "Key not found in %s!\n", keyFile.toLatin1().constData());
			}
		} else {
			fprintf(stderr, "Loading key from %s...\n", keyFile.toLatin1().constData());
		}
		encryptor = new EncryptorSalsa20(key);
	}
}

Cipher::EncryptionAlgorithm Cipher::getAlgorithm() const
{
    return algorithm;
}

QString Cipher::getAlgorithmName() const
{
	return encryptor->getAlgorithmName();
}

QByteArray Cipher::getKey() const
{
	QFile file(keyFile);
	if (!file.open(QIODevice::ReadOnly))
		return QByteArray();
	return file.readAll();
}

bool Cipher::saveKey(QByteArray key, QString keyFile)
{
	QFile file(keyFile);
	if (!file.open(QIODevice::WriteOnly | QIODevice::Truncate)) {
		qDebug() << "Could not open file" << keyFile;
		return false;
	}
	file.write(key);
	file.close();
	QFile::setPermissions(keyFile, QFile::Permissions(QFile::ReadOwner));
	return true;
}

QByteArray Cipher::generateKey(int length) const
{
	QFile urandom("/dev/urandom");
	if (!urandom.open(QIODevice::ReadOnly | QIODevice::Unbuffered))
		return false;
	QByteArray key = urandom.read(length);
	ASSERT(key.length() == length);
	ASSERT(saveKey(key, keyFile));
	return key;
}

QByteArray Cipher::generateKey(QByteArray password, int length, QString keyFile)
{
	qDebug() << "Generating key from password in" << keyFile;
	for (int i = 0; i < 1024; i++) {
		QByteArray hash1 = QCryptographicHash::hash(password, QCryptographicHash::Sha1);
		QByteArray hash2 = QCryptographicHash::hash(password, QCryptographicHash::Md5);
		QByteArray hash3 = QCryptographicHash::hash(password, QCryptographicHash::Md4);
		password = hash1 + hash2 + hash3;
	}
	QByteArray iv;
	iv.fill(0, 64);
	QScopedPointer<Encryptor> encryptor(new EncryptorSalsa20(password, iv));
	QByteArray salt;
	salt.fill(42, length);
	QByteArray key;
	encryptor->encrypt(salt, key);
	ASSERT(saveKey(key, keyFile));
	return key;
}

qint64 Cipher::blockSize() const
{
    return encryptor->blockSize();
}

void Cipher::encrypt(const QByteArray &clearText, QByteArray &cipherText)
{
    encryptor->encrypt(clearText, cipherText);
}

void Cipher::decrypt(const QByteArray &cipherText, QByteArray &clearText)
{
    encryptor->decrypt(cipherText, clearText);
}

bool Cipher::toStream(QDataStream &s)
{
    qint32 encryptionAlgorithmEncoded = algorithm;
    s << encryptionAlgorithmEncoded;

    if (!encryptor->toStream(s))
        return false;

    return s.status() == QDataStream::Ok;
}

bool Cipher::fromStream(QDataStream &s)
{
    qint32 encryptionAlgorithmEncoded;
    s >> encryptionAlgorithmEncoded;
    if (s.status() != QDataStream::Ok)
        return false;

    if (encryptionAlgorithmEncoded == Cipher::NoEncryption) {
		setAlgorithm(Cipher::NoEncryption, keyFile, false);
        if (!encryptor->fromStream(s))
            return false;
	} else if (encryptionAlgorithmEncoded == Cipher::Aes256) {
		setAlgorithm(Cipher::Aes256, keyFile, false);
		if (!encryptor->fromStream(s))
			return false;
	} else if (encryptionAlgorithmEncoded == Cipher::Salsa20) {
		setAlgorithm(Cipher::Salsa20, keyFile, false);
		if (!encryptor->fromStream(s))
			return false;
	} else {
        return false;
    }

    return true;
}
